from Library import *

#2HDMC_dir = "/home/yancywww/Softare/bin/HiggsRelated/2HDMC-1.8.0_hbmodified"

parser = argparse.ArgumentParser(description='manual to this script')
parser.add_argument('--param-input-file',       type=str,   default = "../output/scan_after.dat")
parser.add_argument('--param-output-file',      type=str,   default = "../output/3_adjust_output.dat") #本次程序的输出文件
parser.add_argument('--further-record-file',    type=str,   default = "../output/4_adjust_record.dat") #运用输出文件后，用single_point程序后结果输出到哪个文件
parser.add_argument('--tuning-param',           type=str,   default = "none") #对文件中的哪一个参数进行调整， none为对每一个都进行微调
parser.add_argument('--tuning-param-downlimit', type=float, default = 0) #参数调整的下限
parser.add_argument('--tuning-param-uplimit',   type=float, default = 100) #参数调整的上限
parser.add_argument('--kind',                   type=str,   default = "tiny") #调整的方法 
# tiny 为对tuning-param设置的参数进行大范围撒点，同时对所有其他参数进行微调
# fix 为对tuning-param设置的参数进行大范围撒点，同时所有其他参数固定
parser.add_argument('--generate-number', type=int, default = 1) #对每一个输入的数据点，变为n个新的点
parser.add_argument('--generate-ratio',  type=float, default = 0.05) #对每一个输入的数据点，变为n个新的点
args = parser.parse_args()

# 读取文件
df_sim=pd.read_csv(args.param_input_file, sep='\s+')


#print("文件一共有 ",len(df_sim), " 行数据。")
#print("数据共有 ", len(df_sim.columns)," 列， 列名为：")
#for i in df_sim.columns:
#    print(i,end=' ')

# 设置参数的撒点方式
df  = df_sim[["mh", "mH", "mA", "mC", "sba", "tb", "L6", "L7", "m12_2", "type"]].copy() 

if args.tuning_param == "none":
    args.kind = "tiny"  #如果什么都不输入， 默认为对每一个参数进行微调
elif args.tuning_param not in df.columns:
    print("输入的tuning_param 不在给定的文件中", tuning_param)
    exit(0)



# 设置输出文件
f=open(args.param_output_file,'w')
str_title="mh         mH         mA         mC         sba         tb         L6         L7         m12_2      type      file\n" 
f.write(str_title)

# 对所有点循环
varu = args.tuning_param_uplimit
vard = args.tuning_param_downlimit
vb_range = args.generate_ratio 
for index, row in df.iterrows():
# 对每一个点设置新的值，如果为 fix 方式， 只改变给定的参量
    if args.kind == "fix":
        for i in range(0, args.generate_number, 1):
            row[args.tuning_param] = random.uniform(0,1)*(varu-vard) + vard 
            content="{:.3e}, {:.3e}, {:.3e}, {:.3e}, {:.3e}, {:.3e}, {:.3e}, {:.3e}, {:.3e}, {:.0f},     {file}\n".format(
                    row["mh"], row["mH"], row["mA"], row["mC"], row["sba"], row["tb"], row["L6"], row["L7"], row["m12_2"], row["type"], file=args.further_record_file)
            f.write(content)
    elif args.kind == "tiny":
        for i in range(0, args.generate_number, 1):
            # 这里默认 mH = 125 GeV, 需要根据情况修改 
            row["mh"]    = (random.uniform(0,1)-0.5)* row["mh"]*    vb_range + row["mh"]
            row["mA"]    = (random.uniform(0,1)-0.5)* row["mA"]*    vb_range + row["mA"]
            row["mC"]    = (random.uniform(0,1)-0.5)* row["mC"]*    vb_range + row["mC"]
            row["tb"]    = (random.uniform(0,1)-0.5)* row["tb"]*    vb_range + row["tb"]
            row["sba"]   = (random.uniform(0,1)-0.5)* row["sba"]*   vb_range + row["sba"]
            row["m12_2"] = (random.uniform(0,1)-0.5)* row["m12_2"]* vb_range + row["m12_2"]
            row[args.tuning_param] = random.uniform(0,1)*(varu-vard) + vard 
            content="{:.3e}, {:.3e}, {:.3e}, {:.3e}, {:.3e}, {:.3e}, {:.3e}, {:.3e}, {:.3e}, {:.0f},     {file}\n".format(
                    row["mh"], row["mH"], row["mA"], row["mC"], row["sba"], row["tb"], row["L6"], row["L7"], row["m12_2"], row["type"], file=args.further_record_file)
            f.write(content)

f.close()
